package com.zhang.third.day07;

import com.zhang.third.utils.Event;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.util.Collector;

/**
 * @title: 实时对账--要求5s内left和right对账成功
 * @author: zhang
 * @date: 2022/4/11 10:29
 */
public class Example1 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<Event> leftDS = env
                .addSource(new SourceFunction<Event>() {
                    @Override
                    public void run(SourceContext<Event> ctx) throws Exception {
                        ctx.collectWithTimestamp(new Event("key-1", "left", 1000L), 1000L);
                        ctx.collectWithTimestamp(new Event("key-2", "left", 4000L), 4000L);
                        ctx.emitWatermark(new Watermark(9000L));
                    }

                    @Override
                    public void cancel() {

                    }
                });

        DataStreamSource<Event> rightDS = env
                .addSource(new SourceFunction<Event>() {
                    @Override
                    public void run(SourceContext<Event> ctx) throws Exception {
                        ctx.collectWithTimestamp(new Event("key-1", "right", 3000L), 1000L);
                        ctx.collectWithTimestamp(new Event("key-3", "right", 7000L), 7000L);
                        ctx.emitWatermark(new Watermark(9000L));
                        ctx.collectWithTimestamp(new Event("key-2", "right", 10000L), 10000L);
                    }

                    @Override
                    public void cancel() {

                    }
                });

        leftDS.keyBy(r -> r.key)
                .connect(rightDS.keyBy(r -> r.key))
                .process(new Match())
                .print();

        env.execute();
    }

    public static class Match extends CoProcessFunction<Event, Event, String> {
        private ValueState<Event> leftState;
        private ValueState<Event> rightState;

        @Override
        public void open(Configuration parameters) throws Exception {
            leftState = getRuntimeContext().getState(
                    new ValueStateDescriptor<Event>(
                            "left-state", Types.POJO(Event.class)
                    )
            );

            rightState = getRuntimeContext().getState(
                    new ValueStateDescriptor<Event>(
                            "right-state", Types.POJO(Event.class)
                    )
            );
        }

        @Override
        public void processElement1(Event value, CoProcessFunction<Event, Event, String>.Context ctx, Collector<String> out) throws Exception {
            // right数据先到达，right一定保存在状态变量
            if (rightState.value() != null) {
                // 输出对账成功
                out.collect(value.key + "对账成功，roght事件先到达");
                // 清空right状态
                rightState.clear();
            }
            // left事件先到达，将自己保存在状态，并注册定时器等待right事件5秒钟
            else {
                leftState.update(value);
                ctx.timerService().registerEventTimeTimer(value.ts + 5000L);
            }
        }

        @Override
        public void processElement2(Event value, CoProcessFunction<Event, Event, String>.Context ctx, Collector<String> out) throws Exception {
            // left数据先到达，left数据一定保存在状态变量里
            if (leftState.value() != null) {
                out.collect(value.key + "对账成功，left事件先到达");
                leftState.clear();
            } else {
                rightState.update(value);
                ctx.timerService().registerEventTimeTimer(value.ts + 5000L);
            }
        }

        @Override
        public void onTimer(long timestamp, CoProcessFunction<Event, Event, String>.OnTimerContext ctx, Collector<String> out) throws Exception {
            if (leftState.value() != null) {
                out.collect(leftState.value().key + "对账失败，right事件没来");
                leftState.clear();
            }
            if (rightState.value() != null) {
                out.collect(rightState.value().key + "对账失败，left事件没来");
                rightState.clear();
            }
        }
    }
}
